"""
模拟旋转电场中离子的摆线运动轨迹

旋转电场中入射的离子受到电场力，在垂直方向上描出摆线
在第二段旋转电场中由于反向相位（180度旋转），描出反向摆线

轨迹方程：
    x = (ZeE)/(mω²) × (1 - cos(ωt))
    y = (ZeE)/(mω²) × (ωt - sin(ωt))
    z = √(2ZeV_acc/m) × t
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import sys
import io

# 设置UTF-8输出（处理Windows控制台编码问题）
if sys.platform == 'win32':
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')

# --- 物理常数 ---
e = 1.602176634e-19  # 基本电荷 [C]
amu_to_kg = 1.660539e-27  # 原子质量单位到千克的转换

# --- 粒子参数（Al+）---
Z = 1  # 电荷数
mass_amu = 26.981539  # 质量 [amu]
m = mass_amu * amu_to_kg  # 质量 [kg]

# --- 电场和电压参数 ---
V_acc = 1000.0  # 加速电压 [V]
E_rf = 10000.0  # RF电场强度 [V/m]（调整到更合理的范围）
f_rf = 10000.0  # RF频率 [Hz]
omega = 2.0 * np.pi * f_rf  # 角频率 [rad/s]

# --- 计算摆线参数 ---
R = Z * e * E_rf / (m * omega**2)  # 摆线特征半径 [m]
v_z = np.sqrt(2 * Z * e * V_acc / m)  # z方向速度 [m/s]

print("=== 摆线运动参数 ===")
print(f"粒子: Al+ (Z={Z}, m={mass_amu:.3f} amu)")
print(f"加速电压: V_acc = {V_acc:.1f} V")
print(f"RF电场强度: E_rf = {E_rf:.2e} V/m")
print(f"RF频率: f = {f_rf:.1f} Hz")
print(f"角频率: ω = {omega:.3e} rad/s")
print(f"摆线特征半径: R = {R*1000:.3f} mm")
print(f"z方向速度: v_z = {v_z:.3e} m/s = {v_z*1e-6:.3f} mm/μs")
print()

# --- 计算不同ωt值的轨迹点（扩展到4π） ---
omega_t_values = [np.pi/2, np.pi, 3*np.pi/2, 2*np.pi, 5*np.pi/2, 3*np.pi, 7*np.pi/2, 4*np.pi]
omega_t_labels = ['π/2', 'π', '3π/2', '2π', '5π/2', '3π', '7π/2', '4π']

print("=== 轨迹计算 ===")
print(f"计算 ωt = {', '.join(omega_t_labels)} 时的轨迹点")
print()

trajectory_data = []
for i, omega_t in enumerate(omega_t_values):
    t = omega_t / omega  # 对应的时间 [s]
    
    # 计算轨迹坐标
    if omega_t <= 2*np.pi:
        # 第一段：0到2π，正向摆线
        x = R * (1 - np.cos(omega_t))  # [m]
        y = R * (omega_t - np.sin(omega_t))  # [m]
    else:
        # 第二段：2π到4π，反向摆线（180度旋转的镜像）
        omega_t_rel = omega_t - 2*np.pi  # 相对相位
        
        # 电场反转180度时的摆线方程
        x_base = R * (1 + np.cos(omega_t_rel))  # 镜像摆线
        y_base = R * (2*np.pi - omega_t_rel + np.sin(omega_t_rel))  # y方向反向运动
        
        # 从第一段结束位置连续（第一段结束时：x=0, y=R*2π）
        x1_end = R * (1 - np.cos(2*np.pi))  # = 0
        # x_base从2R开始（当omega_t_rel=0），需要调整为从0开始
        x = x_base - 2*R + x1_end
        y = y_base  # y从R*2π开始，最终回到0
    
    z = v_z * t  # [m]
    
    trajectory_data.append({
        'omega_t': omega_t,
        'omega_t_label': omega_t_labels[i],
        't': t,
        'x': x,
        'y': y,
        'z': z
    })
    
    print(f"ωt = {omega_t_labels[i]}:")
    print(f"  时间: t = {t*1e6:.3f} μs")
    print(f"  位置: x = {x*1000:.3f} mm, y = {y*1000:.3f} mm, z = {z*1000:.3f} mm")
    print(f"  径向距离: r = {np.sqrt(x**2 + y**2)*1000:.3f} mm")
    print()

# --- 计算连续轨迹（扩展到4π）---
# 第一段：0 到 2π - 正向摆线（相位0）
omega_t_max_1 = 2 * np.pi
omega_t1 = np.linspace(0, omega_t_max_1, 1000)
t1 = omega_t1 / omega

x1 = R * (1 - np.cos(omega_t1))
y1 = R * (omega_t1 - np.sin(omega_t1))
z1 = v_z * t1

# 第二段：2π 到 4π - 反向摆线（相位π，180度旋转）
# 在ωt=2π时，电场相位反转180度，进入第二段
# 当电场反转180度时，摆线应该是第一段的镜像反转
omega_t_max_2 = 4 * np.pi
omega_t2 = np.linspace(2*np.pi, omega_t_max_2, 1000)
t2 = omega_t2 / omega

# 第二段摆线：电场相位反转180度，形成180度旋转的镜像摆线
# 使用相位π的摆线方程：x = R*(1 - cos(ωt + π)), y = R*((ωt + π) - sin(ωt + π))
# 或者，相对于2π：x = R*(1 - cos((ωt-2π) + π)) = R*(1 + cos(ωt-2π))
# 简化：设 omega_t2_rel = omega_t2 - 2π，则：
omega_t2_rel = omega_t2 - 2*np.pi  # 相对于2π的相位，从0到2π

# 电场反转180度时的摆线方程：
# x方向：因为cos(ωt-2π+π) = cos(ωt-2π-π) = -cos(ωt-2π)，所以 1 - cos(ωt-2π+π) = 1 + cos(ωt-2π)
# 但为了形成180度旋转的形状（镜像），应该是：
x2_base = R * (1 + np.cos(omega_t2_rel))  # 注意是+cos，形成镜像
# y方向：由于电场反向，运动也反向，应该从最大值逐渐减小
# 第一段结束时 y = R*2π，第二段应该反向运动
# y = R*2π - R*(omega_t2_rel - sin(omega_t2_rel)) = R*(2π - omega_t2_rel + sin(omega_t2_rel))
y2_base = R * (2*np.pi - omega_t2_rel + np.sin(omega_t2_rel))

# 从第一段结束位置连续过渡
# 第一段结束时：x1[-1] = R*(1-cos(2π)) = 0, y1[-1] = R*(2π - sin(2π)) = R*2π
x1_end = R * (1 - np.cos(2*np.pi))  # = 0
y1_end = R * (2*np.pi - np.sin(2*np.pi))  # = R*2π

# x2_base从2R开始（当omega_t2_rel=0时，x2_base = R*(1+1) = 2R）
# 需要调整使x从0开始：x2 = x2_base - 2R + x1_end = x2_base - 2R
x2 = x2_base - 2*R + x1_end  # x从0开始，最终回到0
y2 = y2_base  # y从R*2π开始，最终回到0（形成180度旋转的镜像）
z2 = v_z * t2

# 合并两段轨迹用于完整绘图
omega_t_full = np.concatenate([omega_t1, omega_t2])
x_full = np.concatenate([x1, x2])
y_full = np.concatenate([y1, y2])
z_full = np.concatenate([z1, z2])

print("=== 第二段反向摆线（2π到4π） ===")
print("在ωt=2π时，电场相位反转180度，轨迹反向")
print(f"第一段结束位置: x = {x1[-1]*1000:.3f} mm, y = {y1[-1]*1000:.3f} mm, z = {z1[-1]*1000:.3f} mm")
print(f"第二段起始位置: x = {x2[0]*1000:.3f} mm, y = {y2[0]*1000:.3f} mm, z = {z2[0]*1000:.3f} mm")
print(f"完整轨迹范围: ωt = 0 到 4π")
print()

# --- 绘制轨迹图 ---
print("=== 绘制轨迹图 ===")

# 创建大图
fig = plt.figure(figsize=(18, 12))

# 1. XY平面投影（摆线形状，完整到4π）
ax1 = fig.add_subplot(221)
# 绘制轨迹线
ax1.plot(x1*1000, y1*1000, 'b-', linewidth=2.5, label='First segment', alpha=0.8)
ax1.plot(x2*1000, y2*1000, 'r-', linewidth=2.5, label='Second segment', alpha=0.8)

# 添加箭头表示方向（很小的箭头）
# 箭头大小与轨迹范围成小比例
trajectory_range = max(np.max(np.abs(x1))*1000, np.max(np.abs(y1))*1000, 
                       np.max(np.abs(x2))*1000, np.max(np.abs(y2))*1000)
# 箭头大小：轨迹范围的0.5-0.8%，很小但可见
head_width_base = max(trajectory_range * 0.005, R*1000 * 0.05, 0.03)
head_length_base = max(trajectory_range * 0.008, R*1000 * 0.08, 0.05)

# 增加箭头数量，使其更明显
# 第一段：添加更多箭头
arrow_indices_1 = np.linspace(len(x1)//6, len(x1)-len(x1)//6, 6, dtype=int)  # 第一段添加6个箭头
for idx in arrow_indices_1:
    if idx < len(x1)-1:
        dx = (x1[idx+1] - x1[idx]) * 1000
        dy = (y1[idx+1] - y1[idx]) * 1000
        # 计算箭头方向长度
        arrow_len = np.sqrt(dx**2 + dy**2)
        if arrow_len > 1e-6:  # 确保有足够的长度
            # 归一化方向向量
            dx_norm = dx / arrow_len * head_length_base * 2
            dy_norm = dy / arrow_len * head_length_base * 2
            ax1.arrow(x1[idx]*1000, y1[idx]*1000, dx_norm, dy_norm,
                     head_width=head_width_base, head_length=head_length_base, 
                     fc='blue', ec='darkblue', alpha=0.9, linewidth=1.5, zorder=10)

# 第二段：添加更多箭头
arrow_indices_2 = np.linspace(len(x2)//6, len(x2)-len(x2)//6, 6, dtype=int)  # 第二段添加6个箭头
for idx in arrow_indices_2:
    if idx < len(x2)-1:
        dx = (x2[idx+1] - x2[idx]) * 1000
        dy = (y2[idx+1] - y2[idx]) * 1000
        arrow_len = np.sqrt(dx**2 + dy**2)
        if arrow_len > 1e-6:
            dx_norm = dx / arrow_len * head_length_base * 2
            dy_norm = dy / arrow_len * head_length_base * 2
            ax1.arrow(x2[idx]*1000, y2[idx]*1000, dx_norm, dy_norm,
                     head_width=head_width_base, head_length=head_length_base, 
                     fc='red', ec='darkred', alpha=0.9, linewidth=1.5, zorder=10)

colors = ['green', 'orange', 'purple', 'brown', 'cyan', 'magenta', 'yellow', 'pink']
for i, data in enumerate(trajectory_data):
    if data['omega_t'] <= 2*np.pi:
        ax1.plot(data['x']*1000, data['y']*1000, 'o', color=colors[i % len(colors)], 
                markersize=10, markeredgecolor='k', markeredgewidth=1.5, label='')
        if i < 4:  # 只标注前4个点，避免过于拥挤
            ax1.annotate(f"ωt={data['omega_t_label']}", 
                        xy=(data['x']*1000, data['y']*1000),
                        xytext=(5, 5), textcoords='offset points', fontsize=9,
                        bbox=dict(boxstyle='round,pad=0.3', facecolor='yellow', alpha=0.7))
    else:
        # 第二段的点（2π到4π，不添加到legend）
        ax1.plot(data['x']*1000, data['y']*1000, 'o', color=colors[i % len(colors)], 
                markersize=10, markeredgecolor='k', markeredgewidth=1.5, label='')
        if i >= 4 and i < 7:  # 标注第二段的前几个点
            ax1.annotate(f"ωt={data['omega_t_label']}", 
                        xy=(data['x']*1000, data['y']*1000),
                        xytext=(5, 5), textcoords='offset points', fontsize=9,
                        bbox=dict(boxstyle='round,pad=0.3', facecolor='lightblue', alpha=0.7))

ax1.plot(0, 0, 'ko', markersize=12, markeredgecolor='k', markeredgewidth=2, label='Start')
ax1.axhline(y=0, color='k', linestyle='--', linewidth=0.5, alpha=0.3)
ax1.axvline(x=0, color='k', linestyle='--', linewidth=0.5, alpha=0.3)
ax1.set_xlabel('x [mm]', fontsize=12, fontweight='bold')
ax1.set_ylabel('y [mm]', fontsize=12, fontweight='bold')
ax1.set_title('(a) XY Plane Projection (Cycloid)', fontsize=14, fontweight='bold')
#ax1.legend(fontsize=9, loc='best')
ax1.grid(True, alpha=0.3)
ax1.set_aspect('equal', adjustable='box')

# 2. XZ平面投影
ax2 = fig.add_subplot(222)
ax2.plot(z1*1000, x1*1000, 'b-', linewidth=2.5, label='First segment', alpha=0.8)
ax2.plot(z2*1000, x2*1000, 'r-', linewidth=2.5, label='Second segment', alpha=0.8)

# 添加箭头（很小的箭头）
trajectory_range_xz = max(np.max(np.abs(z1))*1000, np.max(np.abs(x1))*1000, 
                          np.max(np.abs(z2))*1000, np.max(np.abs(x2))*1000)
head_width_xz = max(trajectory_range_xz * 0.005, R*1000 * 0.05, 0.03)
head_length_xz = max(trajectory_range_xz * 0.008, R*1000 * 0.08, 0.05)

arrow_indices_1_xz = np.linspace(len(z1)//6, len(z1)-len(z1)//6, 6, dtype=int)
for idx in arrow_indices_1_xz:
    if idx < len(z1)-1 and idx < len(x1)-1:
        dz = (z1[idx+1] - z1[idx]) * 1000
        dx = (x1[idx+1] - x1[idx]) * 1000
        arrow_len = np.sqrt(dz**2 + dx**2)
        if arrow_len > 1e-6:
            dz_norm = dz / arrow_len * head_length_xz * 2
            dx_norm = dx / arrow_len * head_length_xz * 2
            ax2.arrow(z1[idx]*1000, x1[idx]*1000, dz_norm, dx_norm,
                     head_width=head_width_xz, head_length=head_length_xz, 
                     fc='blue', ec='darkblue', alpha=0.9, linewidth=1.5, zorder=10)

arrow_indices_2_xz = np.linspace(len(z2)//6, len(z2)-len(z2)//6, 6, dtype=int)
for idx in arrow_indices_2_xz:
    if idx < len(z2)-1 and idx < len(x2)-1:
        dz = (z2[idx+1] - z2[idx]) * 1000
        dx = (x2[idx+1] - x2[idx]) * 1000
        arrow_len = np.sqrt(dz**2 + dx**2)
        if arrow_len > 1e-6:
            dz_norm = dz / arrow_len * head_length_xz * 2
            dx_norm = dx / arrow_len * head_length_xz * 2
            ax2.arrow(z2[idx]*1000, x2[idx]*1000, dz_norm, dx_norm,
                     head_width=head_width_xz, head_length=head_length_xz, 
                     fc='red', ec='darkred', alpha=0.9, linewidth=1.5, zorder=10)

# 标记关键点
for i, data in enumerate(trajectory_data):
    ax2.plot(data['z']*1000, data['x']*1000, 'o', color=colors[i % len(colors)], 
            markersize=10, markeredgecolor='k', markeredgewidth=1.5)

ax2.plot(0, 0, 'ko', markersize=12, markeredgecolor='k', markeredgewidth=2, label='Start')
ax2.axhline(y=0, color='k', linestyle='--', linewidth=0.5, alpha=0.3)
ax2.set_xlabel('z [mm]', fontsize=12, fontweight='bold')
ax2.set_ylabel('x [mm]', fontsize=12, fontweight='bold')
ax2.set_title('(b) XZ Plane Projection', fontsize=14, fontweight='bold')
#ax2.legend(fontsize=9, loc='best')
ax2.grid(True, alpha=0.3)

# 3. YZ平面投影
ax3 = fig.add_subplot(223)
ax3.plot(z1*1000, y1*1000, 'b-', linewidth=2.5, label='First segment', alpha=0.8)
ax3.plot(z2*1000, y2*1000, 'r-', linewidth=2.5, label='Second segment', alpha=0.8)

# 添加箭头（很小的箭头）
trajectory_range_yz = max(np.max(np.abs(z1))*1000, np.max(np.abs(y1))*1000,
                          np.max(np.abs(z2))*1000, np.max(np.abs(y2))*1000)
head_width_yz = max(trajectory_range_yz * 0.005, R*1000 * 0.05, 0.03)
head_length_yz = max(trajectory_range_yz * 0.008, R*1000 * 0.08, 0.05)

arrow_indices_1_yz = np.linspace(len(z1)//6, len(z1)-len(z1)//6, 6, dtype=int)
for idx in arrow_indices_1_yz:
    if idx < len(z1)-1 and idx < len(y1)-1:
        dz = (z1[idx+1] - z1[idx]) * 1000
        dy = (y1[idx+1] - y1[idx]) * 1000
        arrow_len = np.sqrt(dz**2 + dy**2)
        if arrow_len > 1e-6:
            dz_norm = dz / arrow_len * head_length_yz * 2
            dy_norm = dy / arrow_len * head_length_yz * 2
            ax3.arrow(z1[idx]*1000, y1[idx]*1000, dz_norm, dy_norm,
                     head_width=head_width_yz, head_length=head_length_yz, 
                     fc='blue', ec='darkblue', alpha=0.9, linewidth=1.5, zorder=10)

arrow_indices_2_yz = np.linspace(len(z2)//6, len(z2)-len(z2)//6, 6, dtype=int)
for idx in arrow_indices_2_yz:
    if idx < len(z2)-1 and idx < len(y2)-1:
        dz = (z2[idx+1] - z2[idx]) * 1000
        dy = (y2[idx+1] - y2[idx]) * 1000
        arrow_len = np.sqrt(dz**2 + dy**2)
        if arrow_len > 1e-6:
            dz_norm = dz / arrow_len * head_length_yz * 2
            dy_norm = dy / arrow_len * head_length_yz * 2
            ax3.arrow(z2[idx]*1000, y2[idx]*1000, dz_norm, dy_norm,
                     head_width=head_width_yz, head_length=head_length_yz, 
                     fc='red', ec='darkred', alpha=0.9, linewidth=1.5, zorder=10)

# 标记关键点
for i, data in enumerate(trajectory_data):
    ax3.plot(data['z']*1000, data['y']*1000, 'o', color=colors[i % len(colors)], 
            markersize=10, markeredgecolor='k', markeredgewidth=1.5)

ax3.plot(0, 0, 'ko', markersize=12, markeredgecolor='k', markeredgewidth=2, label='Start')
ax3.axhline(y=0, color='k', linestyle='--', linewidth=0.5, alpha=0.3)
ax3.set_xlabel('z [mm]', fontsize=12, fontweight='bold')
ax3.set_ylabel('y [mm]', fontsize=12, fontweight='bold')
ax3.set_title('(c) YZ Plane Projection', fontsize=14, fontweight='bold')
#ax3.legend(fontsize=9, loc='best')
ax3.grid(True, alpha=0.3)

# 4. 3D轨迹
ax4 = fig.add_subplot(224, projection='3d')
ax4.plot(x1*1000, y1*1000, z1*1000, 'b-', linewidth=2.5, label='First segment', alpha=0.8)
ax4.plot(x2*1000, y2*1000, z2*1000, 'r-', linewidth=2.5, label='Second segment', alpha=0.8)

# 标记关键点
for i, data in enumerate(trajectory_data):
    ax4.plot([data['x']*1000], [data['y']*1000], [data['z']*1000], 
            'o', color=colors[i % len(colors)], markersize=10, markeredgecolor='k', markeredgewidth=1.5)

ax4.plot([0], [0], [0], 'ko', markersize=12, markeredgecolor='k', markeredgewidth=2, label='Start')
ax4.set_xlabel('x [mm]', fontsize=12, fontweight='bold')
ax4.set_ylabel('y [mm]', fontsize=12, fontweight='bold')
ax4.set_zlabel('z [mm]', fontsize=12, fontweight='bold')
ax4.set_title('(d) 3D Cycloid Trajectory', fontsize=14, fontweight='bold')
#ax4.legend(fontsize=9)

plt.suptitle(f'Cycloid Motion in Rotating Electric Field (0 to 4π)\n', 
             fontsize=14, fontweight='bold', y=0.995)

# 添加统一的legend（在底部）
# 只从第一个子图收集legend，确保每个项目只显示一次
handles1, labels1 = ax1.get_legend_handles_labels()

# 由于所有子图使用相同的label，只需要ax1的legend即可
all_handles = handles1
all_labels = labels1

# 在figure底部添加统一的legend
fig.legend(all_handles, all_labels, loc='lower center', ncol=len(all_handles), 
          fontsize=10, framealpha=0.9, bbox_to_anchor=(0.5, 0.02))

plt.tight_layout(rect=[0, 0.05, 1, 0.97])  # 为底部legend留出空间
plt.savefig('result/cycloid_trajectory.png', dpi=300, bbox_inches='tight')
print("轨迹图已保存: cycloid_trajectory.png")
plt.close()

# --- 绘制单独的XY平面图（更清晰，显示到4π）---
print("绘制XY平面摆线详细图（0到4π）...")
fig_xy = plt.figure(figsize=(14, 12))

# 更密集的轨迹点用于显示摆线的细节
omega_t1_dense = np.linspace(0, 2*np.pi, 1500)  # 第一段：0到2π
omega_t2_dense = np.linspace(2*np.pi, 4*np.pi, 1500)  # 第二段：2π到4π
omega_t_dense = np.concatenate([omega_t1_dense, omega_t2_dense])

# 第一段：0到2π
x1_dense = R * (1 - np.cos(omega_t1_dense))
y1_dense = R * (omega_t1_dense - np.sin(omega_t1_dense))

# 第二段：从2π开始，反向（180度旋转的镜像）
omega_t2_rel_dense = omega_t2_dense - 2*np.pi  # 相对相位
x2_base_dense = R * (1 + np.cos(omega_t2_rel_dense))  # 镜像摆线
y2_base_dense = R * (2*np.pi - omega_t2_rel_dense + np.sin(omega_t2_rel_dense))  # y反向
# 从第一段结束位置连续
x1_end_dense = x1_dense[-1]  # = 0
x2_dense = x2_base_dense - 2*R + x1_end_dense  # x从0开始
y2_dense = y2_base_dense  # y从R*2π开始，最终回到0

ax_xy = fig_xy.add_subplot(111)
ax_xy.plot(x1_dense*1000, y1_dense*1000, 'b-', linewidth=3, label='First segment (0-2π, φ=0)', alpha=0.8)
ax_xy.plot(x2_dense*1000, y2_dense*1000, 'r-', linewidth=3, label='Second segment (2π-4π, φ=π, 180° rotated)', alpha=0.8)

# 添加箭头表示方向（很小很小的箭头）
trajectory_range_xy = max(np.max(np.abs(x1_dense))*1000, np.max(np.abs(y1_dense))*1000,
                          np.max(np.abs(x2_dense))*1000, np.max(np.abs(y2_dense))*1000)
head_width_xy = max(trajectory_range_xy * 0.01, 0.05)  # 很小的箭头
head_length_xy = max(trajectory_range_xy * 0.015, 0.08)  # 很小的箭头

# 第一段箭头
arrow_num_1 = 8  # 第一段添加8个箭头
arrow_idx_1 = np.linspace(len(x1_dense)//8, len(x1_dense)-len(x1_dense)//8, arrow_num_1, dtype=int)
for idx in arrow_idx_1:
    if idx < len(x1_dense)-1:
        dx = (x1_dense[idx+1] - x1_dense[idx]) * 1000
        dy = (y1_dense[idx+1] - y1_dense[idx]) * 1000
        if abs(dx) > 1e-6 or abs(dy) > 1e-6:
            ax_xy.arrow(x1_dense[idx]*1000, y1_dense[idx]*1000, dx*0.25, dy*0.25,
                       head_width=head_width_xy, head_length=head_length_xy, 
                       fc='blue', ec='darkblue', alpha=0.9, linewidth=1.5)

# 第二段箭头
arrow_num_2 = 8  # 第二段添加8个箭头
arrow_idx_2 = np.linspace(len(x2_dense)//8, len(x2_dense)-len(x2_dense)//8, arrow_num_2, dtype=int)
for idx in arrow_idx_2:
    if idx < len(x2_dense)-1:
        dx = (x2_dense[idx+1] - x2_dense[idx]) * 1000
        dy = (y2_dense[idx+1] - y2_dense[idx]) * 1000
        if abs(dx) > 1e-6 or abs(dy) > 1e-6:
            ax_xy.arrow(x2_dense[idx]*1000, y2_dense[idx]*1000, dx*0.25, dy*0.25,
                       head_width=head_width_xy, head_length=head_length_xy, 
                       fc='red', ec='darkred', alpha=0.9, linewidth=1.5)

# 标记所有关键点（0到4π）
for i, data in enumerate(trajectory_data):
    ax_xy.plot(data['x']*1000, data['y']*1000, 'o', color=colors[i % len(colors)], 
              markersize=12, markeredgecolor='k', markeredgewidth=2,
              zorder=5)
    # 标注所有关键点
    if data['omega_t'] <= 2*np.pi:
        # 第一段的点用黄色标注
        ax_xy.annotate(f"ωt = {data['omega_t_label']}", 
                      xy=(data['x']*1000, data['y']*1000),
                      xytext=(10, 10), textcoords='offset points', fontsize=10,
                      bbox=dict(boxstyle='round,pad=0.5', facecolor='yellow', alpha=0.8),
                      fontweight='bold')
    else:
        # 第二段的点用浅蓝色标注
        ax_xy.annotate(f"ωt = {data['omega_t_label']}", 
                      xy=(data['x']*1000, data['y']*1000),
                      xytext=(10, 10), textcoords='offset points', fontsize=10,
                      bbox=dict(boxstyle='round,pad=0.5', facecolor='lightblue', alpha=0.8),
                      fontweight='bold')

ax_xy.plot(0, 0, 'ko', markersize=14, markeredgecolor='k', markeredgewidth=2, 
          label='Start (0,0)', zorder=6)
ax_xy.axhline(y=0, color='k', linestyle='--', linewidth=0.8, alpha=0.4)
ax_xy.axvline(x=0, color='k', linestyle='--', linewidth=0.8, alpha=0.4)
ax_xy.set_xlabel('x [mm]', fontsize=14, fontweight='bold')
ax_xy.set_ylabel('y [mm]', fontsize=14, fontweight='bold')
ax_xy.set_title('Cycloid Motion in Rotating Electric Field (XY Projection, 0 to 4π)', 
               fontsize=16, fontweight='bold')
ax_xy.legend(fontsize=12, loc='best', framealpha=0.9)
ax_xy.grid(True, alpha=0.3, linestyle='--')
ax_xy.set_aspect('equal', adjustable='box')


plt.tight_layout()
plt.savefig('result/cycloid_trajectory_xy.png', dpi=300, bbox_inches='tight')
print("XY平面摆线图已保存: cycloid_trajectory_xy.png")
plt.close()

print()
print("=== 完成 ===")
print("所有图表已生成并保存")

